# models/price_predictor.py

import torch
import torch.nn as nn

class PricePredictor(nn.Module):
    def __init__(self, num_numeric_features: int, num_film_ids: int,
                 film_emb_dim: int = 16, hidden_dims=[128, 64], dropout=0.2):
        """
        :param num_numeric_features: 数值型特征数量（如8个）
        :param num_film_ids: 电影 ID 的种类数，用于 Embedding
        :param film_emb_dim: 电影 ID 的嵌入维度
        :param hidden_dims: MLP 的隐藏层维度
        :param dropout: Dropout 概率
        """
        super(PricePredictor, self).__init__()

        self.film_embedding = nn.Embedding(num_embeddings=num_film_ids, embedding_dim=film_emb_dim)

        input_dim = num_numeric_features + film_emb_dim

        layers = []
        for hidden_dim in hidden_dims:
            layers.append(nn.Linear(input_dim, hidden_dim))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(dropout))
            input_dim = hidden_dim
        layers.append(nn.Linear(input_dim, 1))  # 输出为票价（回归）

        self.mlp = nn.Sequential(*layers)

    def forward(self, x_numeric, x_film_id):
        """
        :param x_numeric: 标准化数值特征张量, shape = [batch_size, num_features]
        :param x_film_id: 影片 ID 张量, shape = [batch_size]
        :return: 票价预测, shape = [batch_size, 1]
        """
        film_emb = self.film_embedding(x_film_id)  # shape = [batch_size, emb_dim]
        x = torch.cat([x_numeric, film_emb], dim=1)
        out = self.mlp(x)
        return out
